/*
 * @Descripttion: 
 * @version: 
 * @Author: yuxintao
 * @Date: 2022-05-24 15:09:27
 * @LastEditors: yuxintao
 * @LastEditTime: 2022-05-24 21:30:14
 */
#include <linux/netfilter_ipv4.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/skbuff.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/if_ether.h>
#include <linux/if_packet.h>

/*
钩子函数的返回值
#define NF_DROP 0   //丢弃该数据包，释放为他分配的资源
#define NF_ACCEPT 1  //保留该数据包，交由下一个hook函数处理
#define NF_STOLEN 2 //忘掉该数据包，hook函数处理该数据包，不再经过netfilter处理
#define NF_QUEUE 3 //将该数据包插入到用户空间
#define NF_REPEAT 4 //再次调用该hook函数
#define NF_STOP 5 //强于NF_ACCEPT，完全接受该数据包，并且后面的所有hook函数都不需处理该数据包
#define NF_MAX_VERDICT NF_STOP 
*/
/*
钩子函数的优先级
enum nf_ip_hook_priorities {
    NF_IP_PRI_FIRST = INT_MIN,
    NF_IP_PRI_CONNTRACK_DEFRAG = -400,
    NF_IP_PRI_RAW = -300,
    NF_IP_PRI_SELINUX_FIRST = -225,
    NF_IP_PRI_CONNTRACK = -200,
    NF_IP_PRI_MANGLE = -150,
    NF_IP_PRI_NAT_DST = -100,
    NF_IP_PRI_FILTER = 0,
    NF_IP_PRI_SECURITY = 50,
    NF_IP_PRI_NAT_SRC = 100,
    NF_IP_PRI_SELINUX_LAST = 225,
    NF_IP_PRI_CONNTRACK_CONFIRM = INT_MAX,
    NF_IP_PRI_LAST = INT_MAX,
};
*/
/*
钩子的安装点
enum nf_inet_hooks 
{
	NF_INET_PRE_ROUTING,
	NF_INET_LOCAL_IN,
	NF_INET_FORWARD,
	NF_INET_LOCAL_OUT,
	NF_INET_POST_ROUTING,
	NF_INET_NUMHOOKS
};*/

//版权声明
MODULE_LICENSE("Dual BSD/GPL");

#define NF_SUCCESS 0
#define NF_FAILURE 1

//判断是否禁止TCP的端口
#define IS_BANDPORT_TCP(status) (status.band_port.port != 0 && status.band_port.protocol == IPPROTO_TCP)

//判断是否禁止UDP的端口
#define IS_BANDPORT_UDP(status) (status.band_port.port != 0 && status.band_port.protocol == IPPROTO_UDP)

//判断是否禁止 PING
#define IS_BANDPING(status) (status.band_ping)

//判断是否禁止IP协议
#define IS_BANDIP(status) (status.band_ip)

#define SOE_BANDIP 0X128
#define SOE_BANDPORT 0X129
#define SOE_BANDPING 0X12A

typedef struct nf_bandport{
    unsigned short protocol;//2个字节
    unsigned short port;//2个字节
}nf_bandport;

typedef struct band_status{
    unsigned int band_ip;//4个字节 禁止的目的ip地址
    nf_bandport band_port;//禁止的端口(TCP OR UDP)
    unsigned char band_ping;//1个字节 禁止ping回显的标志
}band_status;

//初始化绑定状态
band_status b_status;

/* nf sock 选项扩展操作*/
static int nf_sockopt_set(struct sock *sock, int cmd, sockptr_t u, unsigned int len)
{
	int ret = 0;
	band_status status;
	/* 权限检查 */
	if(!capable(CAP_NET_ADMIN))				/*没有足够权限*/
	{
		ret = -EPERM;
		goto ERROR;
	}
    
	/* 从用户空间复制数据*/
	ret = copy_from_user(&status, u.user, len);
	if(ret != 0)								/*复制数据失败*/
	{
		ret = -EINVAL;
		goto ERROR;
	}
	
	/* 命令类型 */
	switch(cmd)//判断setsockopt是管理以下三种中哪种模式
	{
		case SOE_BANDIP:							/*禁止IP协议*/
			/* 设置禁止IP协议 */
			if(IS_BANDIP(status))					/*设置禁止IP协议*/
				b_status.band_ip = status.band_ip; //内核中保存禁止发送的源ip地址
			else									/*取消禁止*/
				b_status.band_ip = 0;
	
			break;
		case SOE_BANDPORT:						/*禁止端口*/
			/* 设置端口禁止和相关的协议类型 */
			if(IS_BANDPORT_TCP(status))				/*禁止TCP*/
			{
				b_status.band_port.protocol = IPPROTO_TCP; //内核中保存禁止访问的TCP端口
				b_status.band_port.port = status.band_port.port;
			}
			else if(IS_BANDPORT_UDP(status))			/*禁止UDP*/
			{
				b_status.band_port.protocol = IPPROTO_UDP; //内核中保存禁止访问的UDP端口
				b_status.band_port.port = status.band_port.port;
			}
			else									/*其他*/
			{
				b_status.band_port.protocol = 0;
				b_status.band_port.port = 0;
			}
			
			break;
		case SOE_BANDPING:						/*禁止ping*/
			if(IS_BANDPING(status))				/*禁止PING*/
			{
				b_status.band_ping = 1; //内核中设置禁止ping回显标志位1
			}
			else 									/*取消禁止*/
			{
				b_status.band_ping = 0;
			}
			
			break;
		default:									/*其他为错误命令*/
			ret = -EINVAL;
			break;
	}
	
ERROR:
	return ret;
}
/* nf sock 操作扩展命令操作*/
static int nf_sockopt_get(struct sock *sock, int cmd, void* __user user, int* len)
{
	int ret = 0;
	/* 权限检查*/
	if(!capable(CAP_NET_ADMIN))				/*没有权限*/
	{
		ret = -EPERM;
		goto ERROR;
	}	
	/* 将数据从内核空间复制到用户空间 */
	switch(cmd)
	{
		case SOE_BANDIP:
		case SOE_BANDPORT:
		case SOE_BANDPING:
			/*复制数据*/
			ret = copy_to_user(user, &b_status, *len);//直接将内核的b_status拷贝到用户态
			if(ret != 0)								/*复制数据失败*/
			{
				ret = -EINVAL;
				goto ERROR;
			}
			break;
		default:
			ret = -EINVAL;
			break;
	}
	
ERROR:
	return ret;
}

/* 在LOCAL_OUT上挂接钩子 */
static unsigned int nf_hook_out(void* hooknum, struct sk_buff *skb, const struct nf_hook_state *state)
{
	struct sk_buff *sk = skb;
	struct iphdr *iph = ip_hdr(sk);//锁定ip头部
	unsigned int dst_ip = iph->daddr;
	if(IS_BANDIP(b_status))						/*判断是否禁止IP协议*/
	{
		if(b_status.band_ip == iph->daddr)			/*IP地址符合*/
		{
			printk(KERN_ALERT "packet to %u.%u.%u.%u is drop\n", (dst_ip&0xff000000ff)>>0, (dst_ip&0xff0000ff00)>>8, 
																(dst_ip&0x00ff0000)>>16, (dst_ip&0xff000000)>>24);
			return NF_DROP;							/*丢弃该网络报文*/
		}
	}
	return NF_ACCEPT;
}

/* 在LOCAL_IN挂接钩子 */
static unsigned int nf_hook_in(void* hooknum, struct sk_buff *skb, const struct nf_hook_state *state)
{
	struct sk_buff *sk = skb;
	struct iphdr *iph = ip_hdr(sk);//锁定ip头部
	unsigned int src_ip = iph->saddr;//数据包的源ip地址
	struct tcphdr *tcph = NULL;
	struct udphdr *udph = NULL;
	
	switch(iph->protocol)								/*IP协议类型*/
	{
		case IPPROTO_TCP:								/*TCP协议*/
			/*丢弃禁止端口的TCP数据*/
			if(IS_BANDPORT_TCP(b_status))
			{
				tcph = (struct tcphdr *) ((char *) iph + (iph->ihl << 2));					/*获得TCP头*/
				if(tcph->dest == b_status.band_port.port)	/*端口匹配*/
				{
					printk(KERN_ALERT "DROP TCP packet from %u.%u.%u.%u, port = %u\n", (src_ip&0xff000000ff)>>0, (src_ip&0xff0000ff00)>>8, 
																		(src_ip&0x00ff0000)>>16, (src_ip&0xff000000)>>24, tcph->dest);
					return NF_DROP;							/*丢弃该数据*/
				}
			}
			break;
		case IPPROTO_UDP:								/*UDP协议*/
			/*丢弃UDP数据*/
			if(IS_BANDPORT_UDP(b_status))					/*设置了丢弃UDP协议*/
			{
				udph = (struct udphdr *) ((char *) iph + (iph->ihl << 2));							/*UDP头部*/
				if(udph->dest == b_status.band_port.port)	/*UDP端口判定*/
				{
					printk(KERN_ALERT "DROP UDP packet from %u.%u.%u.%u, port = %u\n", (src_ip&0xff000000ff)>>0, (src_ip&0xff0000ff00)>>8, 
																		(src_ip&0x00ff0000)>>16, (src_ip&0xff000000)>>24, udph->dest);
					return NF_DROP;								/*丢弃该数据*/
				}
			}
			break;
		case IPPROTO_ICMP:								/*ICMP协议*/
			/*丢弃ICMP报文*/
			if(IS_BANDPING(b_status))						/*设置了禁止PING操作*/
			{
				printk(KERN_ALERT "DROP ICMP packet from %u.%u.%u.%u\n", (src_ip&0xff000000ff)>>0, (src_ip&0xff0000ff00)>>8, 
																		(src_ip&0x00ff0000)>>16, (src_ip&0xff000000)>>24);
				return NF_DROP;								/*丢弃该报文*/
			}
			break;
		default:
			break;
	}
	
	return NF_ACCEPT;
}


/* 初始化nfin钩子，在钩子LOCAL_IN上 */
static struct nf_hook_ops nfin = 
{
	.hook = nf_hook_in,//钩子处理函数
	.hooknum = NF_INET_LOCAL_IN,//钩子的位置值
	.pf = PF_INET,//钩子协议族
	.priority = NF_IP_PRI_FIRST//钩子优先级
};
/*初始化nfout钩子，在钩子LOCAL_OUT上*/
static struct nf_hook_ops nfout=
{
	.hook = nf_hook_out,//钩子处理函数
	.hooknum = NF_INET_LOCAL_OUT,//钩子的位置值
	.pf = PF_INET,//钩子的协议族
	.priority = NF_IP_PRI_FIRST//钩子的优先级
};
/* 初始化nf套接字选项 */
static struct nf_sockopt_ops nfsockopt = {
    .pf	= PF_INET,
    .set_optmin = SOE_BANDIP,
    .set_optmax = SOE_BANDIP + 3,
    .set	= nf_sockopt_set,
    .get_optmin = SOE_BANDIP,
    .get_optmax = SOE_BANDIP + 3,
    .get	= nf_sockopt_get,
};


/* 初始化模块 */
static int __init init(void)
{
    nf_register_net_hook(&init_net, &nfin);		/*注册LOCAL_IN的钩子*/
    nf_register_net_hook(&init_net, &nfout);	/*注册LOCAL_OUT的钩子*/
	nf_register_sockopt(&nfsockopt);			/*注册扩展套接字选项*/
	printk(KERN_ALERT "netfilter example 2 init successfully\n");
												/*打印信息*/
	return NF_SUCCESS;
}
/* 清理模块 */
static void __exit exit1(void)
{
	nf_unregister_net_hook(&init_net, &nfin);	/*注销LOCAL_IN的钩子*/
	nf_unregister_net_hook(&init_net, &nfout);	/*注销LOCAL_OUT的钩子*/
	nf_unregister_sockopt(&nfsockopt);			/*注销扩展套接字选项*/
	printk(KERN_ALERT "netfilter example 2 clean successfully\n");
}

module_init(init);								/*初始化模块*/
module_exit(exit1);								/*模块退出*/
